{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#! /usr/bin/env python\n",
    "# coding=utf-8\n",
    "\n",
    "import os\n",
    "import time\n",
    "import shutil\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import core.utils as utils\n",
    "from tqdm import tqdm\n",
    "from core.dataset import Dataset\n",
    "from core.yolov3 import YOLOv3, decode, compute_loss\n",
    "from core.config import cfg\n",
    "\n",
    "trainset = Dataset('train')\n",
    "logdir = \"./data/log\"\n",
    "steps_per_epoch = len(trainset)\n",
    "global_steps = tf.Variable(1, trainable=False, dtype=tf.int64)\n",
    "warmup_steps = cfg.TRAIN.WARMUP_EPOCHS * steps_per_epoch\n",
    "total_steps = cfg.TRAIN.EPOCHS * steps_per_epoch\n",
    "\n",
    "input_tensor = tf.keras.layers.Input([608, 608, 3])\n",
    "conv_tensors = YOLOv3(input_tensor)\n",
    "\n",
    "output_tensors = []\n",
    "for i, conv_tensor in enumerate(conv_tensors):\n",
    "    pred_tensor = decode(conv_tensor, i)\n",
    "    output_tensors.append(conv_tensor)\n",
    "    output_tensors.append(pred_tensor)\n",
    "\n",
    "model = tf.keras.Model(input_tensor, output_tensors)\n",
    "optimizer = tf.keras.optimizers.Adam()\n",
    "if os.path.exists(logdir): shutil.rmtree(logdir)\n",
    "writer = tf.summary.create_file_writer(logdir)\n",
    "\n",
    "def train_step(image_data, target):\n",
    "    with tf.GradientTape() as tape:\n",
    "        pred_result = model(image_data, training=True)\n",
    "        giou_loss=conf_loss=prob_loss=0\n",
    "\n",
    "        # optimizing process\n",
    "        for i in range(3):\n",
    "            conv, pred = pred_result[i*2], pred_result[i*2+1]\n",
    "            loss_items = compute_loss(pred, conv, *target[i], i)\n",
    "            giou_loss += loss_items[0]\n",
    "            conf_loss += loss_items[1]\n",
    "            prob_loss += loss_items[2]\n",
    "\n",
    "        total_loss = giou_loss + conf_loss + prob_loss\n",
    "\n",
    "        gradients = tape.gradient(total_loss, model.trainable_variables)\n",
    "        optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
    "        if int(global_steps)%100 == 0:\n",
    "            tf.print(\"=> STEP %4d   lr: %.7f   giou_loss: %4.2f   conf_loss: %4.2f   \"\n",
    "                     \"prob_loss: %4.2f   total_loss: %4.2f\" %(global_steps, optimizer.lr.numpy(),\n",
    "                                                              giou_loss, conf_loss,\n",
    "                                                              prob_loss, total_loss))\n",
    "        # update learning rate\n",
    "        global_steps.assign_add(1)\n",
    "        if global_steps < warmup_steps:\n",
    "            lr = global_steps / warmup_steps *cfg.TRAIN.LR_INIT\n",
    "        else:\n",
    "            lr = cfg.TRAIN.LR_END + 0.5 * (cfg.TRAIN.LR_INIT - cfg.TRAIN.LR_END) * (\n",
    "                (1 + tf.cos((global_steps - warmup_steps) / (total_steps - warmup_steps) * np.pi))\n",
    "            )\n",
    "        optimizer.lr.assign(lr.numpy())\n",
    "\n",
    "        # writing summary data\n",
    "        with writer.as_default():\n",
    "            tf.summary.scalar(\"lr\", optimizer.lr, step=global_steps)\n",
    "            tf.summary.scalar(\"loss/total_loss\", total_loss, step=global_steps)\n",
    "            tf.summary.scalar(\"loss/giou_loss\", giou_loss, step=global_steps)\n",
    "            tf.summary.scalar(\"loss/conf_loss\", conf_loss, step=global_steps)\n",
    "            tf.summary.scalar(\"loss/prob_loss\", prob_loss, step=global_steps)\n",
    "        writer.flush()\n",
    "\n",
    "# model.load_weights(\"./yolov3\")\n",
    "for epoch in range(cfg.TRAIN.EPOCHS):\n",
    "    for image_data, target in trainset:\n",
    "        train_step(image_data, target)\n",
    "    model.save_weights(\"./yolov3\")\n",
    "    print('-------------------->epoch:',epoch)\n",
    "    print('-------------------->epoch:',epoch)\n",
    "    print('-------------------->epoch:',epoch)\n",
    "    print('-------------------->epoch:',epoch)\n",
    "    print('-------------------->epoch:',epoch)\n",
    "    print('-------------------->epoch:',epoch)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
